"""
LSTM无人机轨迹预测项目配置文件
整合了基础版本和增强版本的所有配置参数
"""

import os
import numpy as np
import torch
from typing import Dict, Any, Optional, List, Tuple

# =============================================================================
# 基础配置
# =============================================================================
SEED = 42  # 随机种子
BUFFER_SIZE = 1000

# =============================================================================
# 数据配置
# =============================================================================
# 数据路径配置
DATA_FOLDER_NAME = 'uav_data'
DATA_FILE_NAME = 'wrj6jd_noise.dat'
TEST_DATA_FOLDER_NAME = 'uav_data'
TEST_DATA_FILE_NAME = 'wrj6jd_noise.dat'


# 数据参数配置
NUM_DRONES = 6
TOTAL_TIME_SEC = 120
ORIGINAL_DT_SEC = 0.01
NEW_DT_SEC = 0.1
NUM_COORDINATES = 3

# 计算时间步数
TIME_STEPS_ORIGINAL = int(TOTAL_TIME_SEC / ORIGINAL_DT_SEC)
TIME_STEPS_NEW = int(TOTAL_TIME_SEC / NEW_DT_SEC)

# 序列配置
N_STEPS_IN = 100   # LSTM输入序列的步数 (Encoder Input Length)
M_STEPS_OUT = 100  # LSTM输出预测序列的步数 (Decoder Output Length)

# 特征维度
ENCODER_INPUT_DIM = NUM_COORDINATES + NUM_DRONES
DECODER_INPUT_DIM = NUM_COORDINATES + NUM_DRONES
NUM_FEATURES = ENCODER_INPUT_DIM

# =============================================================================
# 模型配置
# =============================================================================
# 模型配置
LSTM_UNITS = 128
DROPOUT_RATE = 0.2
BATCH_SIZE = 32
EPOCHS = 50
LSTM_LAYERS = 2
BIDIRECTIONAL = True
ATTENTION_HEADS = 4
USE_ATTENTION = False
HIDDEN_DIM = 128
BUFFER_SIZE = 1000

# 设置参数
def set_parameters():
    """设置训练参数"""
    global USE_ENHANCED_MODEL
    USE_ENHANCED_MODEL = True


# 损失函数配置
LOSS_FUNCTION = 'Euclidean'  # 可选: 'MSE', 'Euclidean', 'Huber'
# 欧氏距离损失更适合三维轨迹预测，直接优化轨迹点的几何距离

# =============================================================================
# 训练配置 (通用)
# =============================================================================
LEARNING_RATE = 0.001
WEIGHT_DECAY = 1e-4
EARLY_STOPPING_PATIENCE = 15

# 数据处理配置
SMOOTHING_WINDOW_SIZE = 10

# 噪声参数（用于数据增强）
NOISE_MEAN = 0
NOISE_STD_DEV = 0  # 设置为0，不添加噪声

# =============================================================================
# 预测配置
# =============================================================================
# 预测参数
SLIDING_WINDOW_STEP_SEC = 10  # 预测动画的滑动窗口步长（秒）

# =============================================================================
# 模型保存配置
# =============================================================================
MODEL_SAVE_DIR_NAME = 'model'
MODEL_ROOT = 'model'

# 模型文件命名
def get_model_filenames(model_type: str = 'seq2seq') -> Dict[str, str]:
    """
    获取模型文件名称
    Args:
        model_type: 模型类型 ('seq2seq', 'enhanced')
    Returns:
        包含文件名的字典
    """
    base_name = f"{model_type}_model"
    return {
        'model_state_dict': f"{base_name}.pth",
        'scaler': f"{base_name}_scaler.joblib",
        'config': f"{base_name}_config.json"
    }

# =============================================================================
# 可视化配置
# =============================================================================
# 基础可视化
DRONE_TO_VIZ = 10  # 用于可视化的无人机ID (0-indexed)
DRONE_TO_VIZ_NOISE_COMPARE = 0  # 用于噪声对比的无人机ID
NUM_DRONES_TO_PLOT_RELATIVE = 5  # 相对轨迹图显示的无人机数量
NUM_DRONES_TO_PLOT_ABSOLUTE = 3  # 绝对轨迹图显示的无人机数量

# 动画可视化
DRONE_TO_VIZ_ANIMATION = 0  # 动画中主视角无人机ID
ANIMATION_FPS = 10  # 动画帧率

# 颜色配置
COLORMAP = 'tab10'  # matplotlib颜色映射

def get_colormap_and_colors(num_drones: int = None) -> Tuple[Any, List[str]]:
    """
    获取颜色映射和颜色列表
    Args:
        num_drones: 无人机数量，如果为None则使用NUM_DRONES
    Returns:
        (cmap_instance, colors_list)
    """
    import matplotlib.pyplot as plt
    if num_drones is None:
        num_drones = NUM_DRONES
    cmap = plt.colormaps[COLORMAP]
    colors = [cmap(i) for i in np.linspace(0, 1, num_drones)]
    return cmap, colors

# =============================================================================
# 字体配置
# =============================================================================
# Windows字体路径
FONT_PATHS = [
    'C:/Windows/Fonts/simhei.ttf',     # Windows 黑体
    'C:/Windows/Fonts/msyh.ttc',       # Windows 微软雅黑
    'C:/Windows/Fonts/simsun.ttc',     # Windows 宋体
    'C:/Windows/Fonts/simkai.ttf'      # Windows 楷体
]

# =============================================================================
# 系统配置
# =============================================================================
# PyTorch配置
DEVICE = 'cuda' if torch.cuda.is_available() else 'cpu'
PIN_MEMORY = True  # DataLoader是否使用固定内存

# 训练/预测模式选择
MODE = 'train'  # 'train' 或 'predict'

# =============================================================================
# 路径配置
# =============================================================================
def get_data_path(folder_name: str = None, file_name: str = None) -> str:
    """
    获取数据文件完整路径
    Args:
        folder_name: 文件夹名称
        file_name: 文件名称
    Returns:
        完整路径字符串
    """
    folder = folder_name or DATA_FOLDER_NAME
    file = file_name or DATA_FILE_NAME
    return os.path.join(folder, file)

def get_test_data_path(folder_name: str = None, file_name: str = None) -> str:
    """
    获取测试数据文件完整路径
    Args:
        folder_name: 文件夹名称
        file_name: 文件名称
    Returns:
        完整路径字符串
    """
    folder = folder_name or TEST_DATA_FOLDER_NAME
    file = file_name or TEST_DATA_FILE_NAME
    return os.path.join(folder, file)

def get_model_path(model_dir: str = None, filename: str = None) -> str:
    """
    获取模型文件完整路径
    Args:
        model_dir: 模型目录
        filename: 文件名
    Returns:
        完整路径字符串
    """
    model_dir = model_dir or MODEL_SAVE_DIR_NAME
    os.makedirs(model_dir, exist_ok=True)
    return os.path.join(model_dir, filename)

# =============================================================================
# 配置验证
# =============================================================================
def validate_config() -> Dict[str, Any]:
    """
    验证配置参数的有效性
    Returns:
        配置验证结果
    """
    config_issues = []

    # 检查时间步数
    if N_STEPS_IN + M_STEPS_OUT > TIME_STEPS_NEW:
        config_issues.append(f"序列长度 ({N_STEPS_IN + M_STEPS_OUT}) 超过总时间步数 ({TIME_STEPS_NEW})")

    # 检查批次大小
    if BATCH_SIZE <= 0:
        config_issues.append("批次大小必须大于0")

    # 检查学习率
    if LEARNING_RATE <= 0 or LEARNING_RATE > 1:
        config_issues.append("学习率必须在(0, 1]范围内")

    # 检查无人机数量
    if NUM_DRONES <= 0:
        config_issues.append("无人机数量必须大于0")

    return {
        'valid': len(config_issues) == 0,
        'issues': config_issues
    }

# =============================================================================
# 获取配置
# =============================================================================
def get_training_config() -> Dict[str, Any]:
    """
    获取训练配置
    Returns:
        训练配置字典
    """
    return {
        'seed': SEED,
        'epochs': EPOCHS,
        'batch_size': BATCH_SIZE,
        'learning_rate': LEARNING_RATE,
        'weight_decay': WEIGHT_DECAY,
        'early_stopping_patience': EARLY_STOPPING_PATIENCE,
        'device': DEVICE,
        'pin_memory': PIN_MEMORY,
        'training_version': 'enhanced'
    }

def get_prediction_config() -> Dict[str, Any]:
    """
    获取预测配置
    Returns:
        预测配置字典
    """
    return {
        'device': DEVICE,
        'prediction_version': PREDICTION_VERSION,
        'sliding_window_step_sec': SLIDING_WINDOW_STEP_SEC,
        'animation_fps': ANIMATION_FPS
    }

def get_model_config() -> Dict[str, Any]:
    """
    获取模型配置
    Returns:
        模型配置字典
    """
    if USE_ENHANCED_MODEL:
        return {
            'model_type': 'enhanced',
            'lstm_units': LSTM_UNITS,
            'hidden_dim': HIDDEN_DIM,
            'lstm_layers': LSTM_LAYERS,
            'dropout_rate': DROPOUT_RATE,
            'bidirectional': BIDIRECTIONAL,
            'attention_heads': ATTENTION_HEADS,
            'use_attention': USE_ATTENTION,
            'encoder_input_dim': ENCODER_INPUT_DIM,
            'decoder_input_dim': DECODER_INPUT_DIM,
            'num_coordinates': NUM_COORDINATES
        }
    else:
        return {
            'model_type': 'basic',
            'lstm_units': LSTM_UNITS,
            'dropout_rate': DROPOUT_RATE,
            'encoder_input_dim': ENCODER_INPUT_DIM,
            'decoder_input_dim': DECODER_INPUT_DIM,
            'num_coordinates': NUM_COORDINATES
        }

# =============================================================================
# 环境检查
# =============================================================================
def check_environment() -> Dict[str, Any]:
    """
    检查运行环境
    Returns:
        环境信息字典
    """
    import torch
    return {
        'pytorch_version': torch.__version__,
        'cuda_available': torch.cuda.is_available(),
        'cuda_device_count': torch.cuda.device_count() if torch.cuda.is_available() else 0,
        'cuda_device_name': torch.cuda.get_device_name(0) if torch.cuda.is_available() else None,
        'device': DEVICE
    }

# =============================================================================
# Train.py所需配置 (在函数定义之后)
# =============================================================================
# 获取训练数据路径
def get_training_data_path():
    """获取训练数据路径"""
    return get_data_path()

# 获取训练模型文件名
def get_training_model_filenames():
    """获取训练模型文件名 - 基于训练数据文件名"""
    data_file_name = DATA_FILE_NAME.replace('.dat', '')  # 移除文件扩展名
    base_name = f"{data_file_name}_model"
    return {
        'model_state_dict': f"{base_name}.pth",
        'scaler': f"{base_name}_scaler.joblib",
        'config': f"{base_name}_config.json"
    }

# 设置训练路径
TRAIN_DATA_PATH = get_training_data_path()
TEST_DATA_PATH = get_test_data_path()
MODEL_STATE_DICT_PATH = get_model_path(filename=get_training_model_filenames()['model_state_dict'])
SCALER_PATH = get_model_path(filename=get_training_model_filenames()['scaler'])

# 导入时设置参数
set_parameters()

# =============================================================================
# 主配置入口
# =============================================================================
if __name__ == "__main__":
    # 验证配置
    validation = validate_config()
    if not validation['valid']:
        print("配置验证失败:")
        for issue in validation['issues']:
            print(f"  - {issue}")
    else:
        print("配置验证通过")

    # 检查环境
    env_info = check_environment()
    print(f"运行环境: {env_info}")

    # 显示主要配置
    print(f"训练版本: enhanced")
    print(f"预测版本: enhanced")
    print(f"使用增强模型: {USE_ENHANCED_MODEL}")
    print(f"设备: {DEVICE}")
    print(f"数据文件夹: {DATA_FOLDER_NAME}")
    print(f"模型保存目录: {MODEL_SAVE_DIR_NAME}")
    print(f"LSTM单元数: {LSTM_UNITS}")
    print(f"批次大小: {BATCH_SIZE}")
    print(f"训练轮次: {EPOCHS}")
